import torch

class TextClassifier(torch.nn.Module):
    def __init__(self,num_embedding,embedding_dim,padding_idx,cls_num):
        '''
        将输入的整数序列（词的索引）转换为稠密的词向量
        :param num_embedding:词表大小（词汇总数）
        :param embedding_dim:每个词的向量的维度
        :param padding_idx:填充符号的索引
        :param cls_num:类别数量
        '''
        super().__init__()
        self.emb = torch.nn.Embedding(num_embedding,embedding_dim,padding_idx)

        self.fc = torch.nn.Sequential(
            torch.nn.Linear(embedding_dim,128),
            torch.nn.ReLU(),

            torch.nn.Linear(128, 64),
            torch.nn.ReLU(),

            torch.nn.Linear(64, 32),
            torch.nn.ReLU(),

            torch.nn.Linear(32, cls_num),
        )

    def forward(self,x):
        out = self.emb(x)
        out = torch.sum(out,dim=1)
        return self.fc(out)

if __name__ == '__main__':
    # 定义模型参数
    vocab_size = 5000  # 词表大小
    embed_dim = 128  # 词向量维度
    num_classes = 4  # 分类类别数
    padding_idx = 0  # 填充符的索引（假设0是<PAD>）
    net = TextClassifier(vocab_size,embed_dim,padding_idx,num_classes)

    input_batch = torch.LongTensor([
        [1,2,3,4,0,0], #句子1（后两个位置是填充符）
        [5,6,7,8,9,10] #句子2
    ])

    out = net(input_batch)
    print(out.shape)